﻿/*
 *   Copyright 2011 Peter Crossley
 *
 *   Licensed under the Apache License, Version 2.0 (the "License");
 *   you may not use this file except in compliance with the License.
 *   You may obtain a copy of the License at
 *
 *       http://www.apache.org/licenses/LICENSE-2.0
 *
 *   Unless required by applicable law or agreed to in writing, software
 *   distributed under the License is distributed on an "AS IS" BASIS,
 *   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 *   See the License for the specific language governing permissions and
 *   limitations under the License.
 *
 *  ----------------------------------------------------------------------------
 *  File:       SysprepService.cs
 *  Author:     Peter Crossley
 *  ----------------------------------------------------------------------------
 */
using System;
using System.Diagnostics;
using System.IO;
using System.Runtime.InteropServices;
using System.Security.Cryptography;
using System.Management;
using System.DirectoryServices;
using System.Collections;
using System.Collections.Generic;
using System.Xml.Serialization;
using System.Text;
using UEC.Modules.Configuration;
using System.Net;
using System.Windows.Forms;
using NVelocity.App;
using NVelocity;
using NVelocity.Exception;
using Microsoft.Win32;

namespace UEC.Modules.Sysprep
{
    public class SysprepService : ICloudModule
    {
        private CloudManager manager;

        public SysprepService(CloudManager manager)
        {
            this.manager = manager;
            manager.CloudServiceEvent += new CloudManager.CloudServiceHandler(ServiceHandler);
        }

        public static string SYSPREP_SETTINGS
        {
            get { return "SysprepSettings"; }
        }

        #region ICloudModule Members

        public string GetServiceName()
        {
            return global::UEC.Properties.Resources.SysprepServiceName;
        }

        public System.Windows.Forms.UserControl GetConfigurationControl()
        {
            this.sysprepSettings = new SysprepSettings(this);
            return sysprepSettings;
        }

        public decimal GetWeight()
        {
            return global::UEC.Properties.Settings.Default.ServiceWeight;
        }

        public CloudManager Manager
        {
            get { return manager; }
        }

        private SysprepSettings sysprepSettings;
        #endregion

        private void ServiceHandler(object sender, CloudManager.CloudServiceEventArgs e)
        {
            ServiceContext context = manager.ServiceContext;
            manager.LogEvent("ServiceHandler: " + SYSPREP_SETTINGS + " (" + e.CloudServiceState.ToString() + ")", EventLogEntryType.Information,true);
            switch (e.CloudServiceState)
            {
                case CloudManager.CloudServiceStates.SYSPREP_AUDIT:
                    try
                    {
                        context.ServiceBegin();
                        context.Log("Starting sysprep audit.");
                        GeneralizeSysprep();
                    }
                    finally
                    {
                        context.Log("Finished sysprep audit.");
                        context.ServiceEnd();
                    }
                    break;

                case CloudManager.CloudServiceStates.SYSPREP:
                    try
                    {
                        context.ServiceBegin();
                        context.Log("Starting sysprep process.");
                        FinalizeSysprep();
                    }
                    finally
                    {
                        context.Log("Finished sysprep process.");
                        context.ServiceEnd();
                    }
                    break;

            }
            context.FlushContext();
        }


        private void ProcessTemplate(string file)
        {
            if (!File.Exists(file))
            {
                manager.LogEvent("info : no unattended file defined for setup, user interation enabled.", EventLogEntryType.Information, true);
                return;
            }


            /*
	        * setup
	        */
            Velocity.Init();

            /*
            *  Make a context object and populate with the data.  This 
            *  is where the Velocity engine gets the data to resolve the
            *  references (ex. $list) in the template
            */
            VelocityContext context = new VelocityContext();
            foreach (string type in manager.GetTypes())
            {
                context.Put(type, ContextTypeMap(type));
            }
            context.Put("system", manager.SystemProperties);

            try
            {
                TextWriter streamWriter = null;
                string templatePath = Path.Combine(Environment.CurrentDirectory, Path.GetFileName(file) + ".vm");
                try
                {
                    File.Copy(file, templatePath, true);
                    streamWriter = new StreamWriter(file);
                    Velocity.MergeTemplate(Path.GetFileName(templatePath), context, streamWriter);
                    manager.LogEvent("Processed Sysprep template [" + file + "]", EventLogEntryType.Information);
                }
                finally
                {
                    if (streamWriter != null)
                    {
                        streamWriter.Close();
                    }
                    File.Delete(templatePath);
                }
            }
            catch (ResourceNotFoundException)
            {
                manager.LogEvent("error : cannot find template " + file, EventLogEntryType.Error, true);
            }
            catch (ParseErrorException pee)
            {
                manager.LogEvent("Syntax error in template " + file + ":" + pee, EventLogEntryType.Error, true);
            }

            /*
            *  Now have the template engine process your template using the
            *  data placed into the context.  Think of it as a  'merge' 
            *  of the template and the data to produce the output stream.
            */
            
        }

        private IDictionary<string, object> ContextTypeMap(string type)
        {
            IDictionary<string, object> dict = new Dictionary<string, object>();
            foreach (string name in manager.GetValueNames(type)) 
            {
                dict.Add(name, manager.GetValue(type,name, null));
            }
            return dict;
        }

        public void FinalizeSysprep()
        {

            //we are clearing our mode so that we have normal operations again, post sysprep.
            manager.Mode = null;

            try
            {
                //clean up from sysprep OOBE
                RegistryKey key = Registry.LocalMachine;
                key = key.OpenSubKey("Software").OpenSubKey("Microsoft").OpenSubKey("Windows").OpenSubKey("CurrentVersion").OpenSubKey("RunOnce", true);
                if (key != null)
                {
                    foreach (string name in key.GetValueNames())
                    {
                        object data = key.GetValue(name);
                        if (data != null && typeof(string) == data.GetType())
                        {
                            //check to see if we are in OOBE mode still if so clean that up.
                            if (((string)data).ToLower().Contains("oobe.exe"))
                            {
                                key.DeleteValue(name);
                                break;
                            }
                        }
                    }
                }

                //clean out any remain issues from Setup that OOBE missed due to reboot.
                key = Registry.LocalMachine.OpenSubKey("SYSTEM").OpenSubKey("Setup", true);
                key.SetValue("CmdLine", "");
                key.SetValue("OOBEInProgress", 0);
                key.SetValue("RestartSetup", 0);
                key.SetValue("SetupPhase", 0);
                key.SetValue("SetupType", 0);
                key.SetValue("SystemSetupInProgress", 0);

            }
            catch (Exception e) {
                manager.LogEvent("error : cannot clean up from OOBE:\n" + e.ToString(), EventLogEntryType.Error, true);
            }
        }

        public void AuditSysprep()
        {
            //next time we start we will go in to audit mode and finish pulling instance information
            manager.Mode = "audit"; 
            try
            {
                Cursor.Current = Cursors.WaitCursor;
                Process p = new Process();
                p.StartInfo.UseShellExecute = false;
                p.StartInfo.FileName = Path.Combine(Path.Combine(Environment.SystemDirectory, "sysprep"), "sysprep.exe");
                p.StartInfo.Arguments = "/audit /quit";
                p.StartInfo.WorkingDirectory = Path.Combine(Environment.SystemDirectory, "sysprep");
                p.Start();
                p.WaitForExit();

                //setup runonce for audit mode
                RegistryKey key = Registry.LocalMachine;
                key = key.OpenSubKey("Software").OpenSubKey("Microsoft").OpenSubKey("Windows").OpenSubKey("CurrentVersion", true).CreateSubKey("RunOnce");
                key.SetValue("CIPSAuditMode", Application.ExecutablePath + " /audit" + (manager.Debug ? " /debug" : ""), RegistryValueKind.String);
                
                Cursor.Current = Cursors.Default;
                MessageBox.Show(Properties.Resources.SysprepNextSteps, Properties.Resources.SystemMessageTitle);

            }
            catch (Exception e)
            {
                manager.LogEvent("Error audit prep system:\n" + e, EventLogEntryType.Error, true);
                Cursor.Current = Cursors.Default;
                MessageBox.Show(Properties.Resources.ErrorGeneralizingSystem + "\n" + e.Message, Properties.Resources.SystemMessageTitle);
            }


        }

        public void GeneralizeSysprep()
        {
            string value = (string)manager.GetValue(SysprepService.SYSPREP_SETTINGS, "TemplateLocation", "");

            if (!string.IsNullOrEmpty(value))
            {
                string sysprepXml = Path.Combine(Path.Combine(Environment.SystemDirectory, "sysprep"), "unattend.xml");
                //make sure we dont have a temp file
                File.Delete(sysprepXml);

                if (Uri.IsWellFormedUriString(value, UriKind.Absolute))
                {
                    GetRemoteTemplate(value, sysprepXml);
                }
                else
                {
                    File.Copy(value, sysprepXml, true);
                }

                ProcessTemplate(sysprepXml);

                try
                {
                    Process p = new Process();
                    p.StartInfo.UseShellExecute = false;
                    p.StartInfo.FileName = Path.Combine(Path.Combine(Environment.SystemDirectory, "sysprep"), "sysprep.exe");
                    p.StartInfo.Arguments = "/generalize /oobe /unattend:unattend.xml /quit";
                    p.StartInfo.WorkingDirectory = Path.Combine(Environment.SystemDirectory, "sysprep");
                    p.Start();
                    
                    p.WaitForExit();

                    //we are in the sysprep phase
                    manager.Mode = "sysprep";
                   
                    //reboot system
                    manager.Reboot = true;
                }
                catch (Exception e)
                {
                    manager.LogEvent("Error generalizing system:\n" + e, EventLogEntryType.Error, true);
                }
                finally
                {
                    //delete the temp file
                    File.Delete(sysprepXml);
                }
            }
            else
            {
                manager.LogEvent(Properties.Resources.InvalidSysprepTemplate, EventLogEntryType.Error, true);
                return;
            }

            
        }

        private void GetRemoteTemplate(string location, string destinationFile)
        {
            try
            {            
                Uri serviceUri = new Uri(location);
                WebClient downloader = new WebClient();
                StreamWriter writer = null;
                try
                {
                    string template = downloader.DownloadString(serviceUri);
                    if (String.IsNullOrEmpty(template))
                    {
                        //dont have the template
                        return;
                    }
                    writer = File.CreateText(destinationFile);
                    writer.Write(template);
                }
                finally
                {
                    if (writer != null)
                    {
                        writer.Close();
                    }
                }
            }
            catch (Exception e)
            {
                manager.LogEvent("Unable to load remote template:\n" + e.ToString(), EventLogEntryType.Error);
                try
                {
                    File.Delete(destinationFile);
                }
                catch (Exception e1)
                {
                    manager.LogEvent("Unable to delete temporary template:\n" + e1.ToString(), EventLogEntryType.Error);
                }
            }
        }

        
    }
}
